import tkinter as tk
from tkinter import filedialog
import re
import requests
import os
import concurrent.futures

class M3U8DownloaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title('M3U8下载器')

        # URL输入
        tk.Label(root, text='M3U8 URL:').grid(row=0, column=0, padx=5, pady=5)
        self.url_entry = tk.Entry(root, width=60, font=('Arial', 12))
        self.url_entry.grid(row=0, column=1, padx=5, pady=5)

        # 保存目录选择
        tk.Label(root, text='保存目录:').grid(row=1, column=0, padx=5, pady=5)
        self.dir_entry = tk.Entry(root, width=60, font=('Arial', 12))
        self.dir_entry.grid(row=1, column=1, padx=5, pady=5)
        tk.Button(root, text='选择目录', command=self.choose_directory).grid(row=1, column=2, padx=5, pady=5)

        # 下载按钮
        tk.Button(root, text='开始下载', command=self.start_download).grid(row=2, column=1, padx=5, pady=5)

        # 进度显示
        self.progress_label = tk.Label(root, text='')
        self.progress_label.grid(row=3, column=1, padx=5, pady=5)
        
        # 日志显示
        tk.Label(root, text='下载日志:').grid(row=4, column=0, padx=5, pady=5, sticky='nw')
        self.log_text = tk.Text(root, height=15, width=80, state='disabled')
        self.log_text.grid(row=4, column=1, columnspan=2, padx=5, pady=5, sticky='nsew')
        scrollbar = tk.Scrollbar(root, command=self.log_text.yview)
        scrollbar.grid(row=4, column=3, sticky='ns')
        self.log_text.config(yscrollcommand=scrollbar.set)
        
        # 配置网格布局权重
        root.grid_rowconfigure(4, weight=1)
        root.grid_columnconfigure(1, weight=1)

    def choose_directory(self):
        directory = filedialog.askdirectory()
        if directory:
            self.dir_entry.delete(0, tk.END)
            self.dir_entry.insert(0, directory)

    def start_download(self):
        m3u8_url = self.url_entry.get()
        output_dir = self.dir_entry.get()
        if not m3u8_url or not output_dir:
            return

        self.progress_label.config(text='开始下载...')
        self.root.update()

        download_ts_files(m3u8_url, output_dir, self.update_progress)
        self.progress_label.config(text='下载完成！')

    def update_progress(self, progress):
        self.progress_label.config(text=progress)
        self.log_message(progress)
        self.root.update()
        
    def log_message(self, message):
        self.log_text.config(state='normal')
        self.log_text.insert('end', message + '\n')
        self.log_text.config(state='disabled')
        self.log_text.see('end')


def download_ts_files(m3u8_url, output_dir, progress_callback=None):
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 检查m3u8文件是否存在
    m3u8_path = os.path.join(output_dir, 'index.m3u8')
    if os.path.exists(m3u8_path):
        with open(m3u8_path, 'r') as f:
            m3u8_content = f.read()
    else:
        # 下载m3u8文件
        response = requests.get(m3u8_url)
        m3u8_content = response.text
        with open(m3u8_path, 'w') as f:
            f.write(m3u8_content)

    # 解析m3u8文件，获取.ts文件URL
    ts_prefix = '/'.join(m3u8_url.split('/')[:-1])
    # 按行解析m3u8文件，正确处理换行符
    ts_urls = []
    for line in m3u8_content.splitlines():
        line = line.strip()
        if line.endswith('.ts') and not line.startswith('#'):
            ts_urls.append(f'{ts_prefix}/{line}')

    # 使用线程池下载.ts文件
    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
        future_to_url = {executor.submit(download_single_ts, ts_url, output_dir): ts_url for ts_url in ts_urls}
        for future in concurrent.futures.as_completed(future_to_url):
            try:
                future.result()
            except Exception as exc:
                ex_msg = f'下载过程中发生错误: {exc}'
                if progress_callback:
                    progress_callback(ex_msg)

            # 显示下载进度
            downloaded_count = len([f for f in os.listdir(output_dir) if f.endswith('.ts') and not f.startswith('.')])
            total_count = len(ts_urls)
            progress = f'下载进度: {downloaded_count}/{total_count} ({downloaded_count/total_count:.2%})'
            if progress_callback:
                progress_callback(progress)


def download_single_ts(ts_url, output_dir, progress_callback=None):
    file_name = ts_url.split('/')[-1]
    file_path = os.path.join(output_dir, file_name)

    # 检查文件是否存在
    if os.path.exists(file_path):
        if progress_callback:
            progress_callback(f'文件已存在，跳过下载: {file_name}')
        return

    # 下载文件，最多重试3次
    max_retries = 3
    for attempt in range(max_retries):
        try:
            if progress_callback:
                progress_callback(f'正在下载: {file_name} (尝试 {attempt + 1}/{max_retries})')
            response = requests.get(ts_url)
            with open(file_path, 'wb') as f:
                f.write(response.content)
            break
        except Exception as e:
            error_msg = f'下载失败: {file_name}, 错误: {str(e)}'
            if progress_callback:
                progress_callback(error_msg)
            print(error_msg)
            if attempt == max_retries - 1:
                fail_msg = f'文件 {file_name} 下载失败，已达到最大重试次数'
                if progress_callback:
                    progress_callback(fail_msg)
                print(fail_msg)


if __name__ == '__main__':
    root = tk.Tk()
    app = M3U8DownloaderApp(root)
    root.mainloop()